import torch
import torchaudio
import os
from infer.inference import AbsInference


class Inference(AbsInference):
    def __init__(self, h):
        super().__init__(h)

    @torch.no_grad()
    def infer(self, data, idx):
        audio, path = data
        audio = audio.cuda()
        audio = audio.unsqueeze(0)
        emb = self.model(audio, mask_ratio = None)[0] # [B, T, E]
        audio = self.model.recon(emb)  # [B, T]
        base_name = os.path.basename(path)
        base_name = base_name.replace(".wav", "")
        torchaudio.save(f"{self.output_dir}/{base_name}.wav", audio.cpu(), 16000)
        